package com.zyc.utils.threadpool;

import com.google.common.util.concurrent.ThreadFactoryBuilder;
import lombok.extern.slf4j.Slf4j;

import java.util.Map;
import java.util.concurrent.*;

/**
 * 创建不同线程池的工具类
 * @author zyc
 * @create 2022--05--16 16:05
 */
@Slf4j
public class RpcThreadPoolFactoryUtil {
    /**
     * 通过 threadNamePrefix 来区分不同线程池（我们可以把相同 threadNamePrefix 的线程池看作是为同一业务场景服务）。
     * key: threadNamePrefix
     * value: threadPool
     */
    private static final Map<String, ThreadPoolExecutor> POOL_MAP = new ConcurrentHashMap<>();

    public static ThreadPoolExecutor createThreadPoolIfAbsent(String threadNamePrefix) {
        return createThreadPoolIfAbsent(threadNamePrefix, new RpcThreadPoolConfig(), false);
    }

    public static ThreadPoolExecutor createThreadPoolIfAbsent(String threadNamePrefix, RpcThreadPoolConfig rpcThreadPoolConfig) {
        return createThreadPoolIfAbsent(threadNamePrefix, rpcThreadPoolConfig, false);
    }

    public static ThreadPoolExecutor createThreadPoolIfAbsent(String threadNamePrefix, RpcThreadPoolConfig rpcThreadPoolConfig, Boolean daemon) {
        return POOL_MAP.computeIfAbsent(threadNamePrefix, k -> createThreadPool(threadNamePrefix, daemon, rpcThreadPoolConfig));
    }

    /**
     * 创建线程池
     * @param threadNamePrefix
     * @param daemon
     * @param threadPoolConfig
     * @return
     */
    private static ThreadPoolExecutor createThreadPool(String threadNamePrefix, Boolean daemon, RpcThreadPoolConfig threadPoolConfig) {
        ThreadFactory threadFactory = createThreadFactory(threadNamePrefix, daemon);
        return new ThreadPoolExecutor(
                threadPoolConfig.getCorePoolSize(),
                threadPoolConfig.getMaximumPoolSize(),
                threadPoolConfig.getKeepAliveTime(),
                threadPoolConfig.getUnit(),
                threadPoolConfig.getWorkQueue(),
                threadFactory
        );
    }

    /**
     * 创建线程工厂
     * @param threadNamePrefix
     * @param daemon
     * @return
     */
    public static ThreadFactory createThreadFactory(String threadNamePrefix, Boolean daemon) {
        if (threadNamePrefix == null) {
            return Executors.defaultThreadFactory();
        }
        if (daemon == null) {
            return new ThreadFactoryBuilder().setNameFormat(threadNamePrefix + "-%d").build();
        }
        return new ThreadFactoryBuilder().setNameFormat(threadNamePrefix + "-%d").setDaemon(daemon).build();
    }

    public static void shutDownAllThreadPool() {
        log.info("关闭所有线程");
        for (Map.Entry<String, ThreadPoolExecutor> entry : POOL_MAP.entrySet()) {
            ThreadPoolExecutor pool = entry.getValue();
            pool.shutdown();
            log.info("关闭线程池【{}】【{}】", entry.getKey(), pool.isTerminated());
            try {
                pool.awaitTermination(10, TimeUnit.SECONDS);
            } catch (InterruptedException e) {
                log.info("线程池{}没有被终止",entry.getKey());
                e.printStackTrace();
            }
        }
    }
}
